from environment.uav import UAV
from environment.target import Target
from environment.multi_target import MultiTarget
import numpy as np
import math
import random

class EncirclementEnv:
    def __init__(self, num_uavs, grid_size, num_position_action, num_energy_action, seed=None):
        # 如果提供了种子，设置环境的随机种子
        if seed is not None:
            np.random.seed(seed)
            random.seed(seed)
            
        self.grid_size = grid_size
        self.num_uavs = num_uavs
        
        # 创建无人机，初始位置分散在边界
        self.uavs = []
        for i in range(num_uavs):
            uav = UAV(grid_size, i, [], num_position_action, num_energy_action)
            # 将5个无人机初始化在网格边界的不同位置
            if i == 0:
                uav.position = (10, 10)  # 左上
            elif i == 1:
                uav.position = (grid_size-10, 10)  # 右上
            elif i == 2:
                uav.position = (10, grid_size-10)  # 左下
            elif i == 3:
                uav.position = (grid_size-10, grid_size-10)  # 右下
            elif i == 4:
                uav.position = (grid_size//2, 10)  # 上中
            else:
                # 更多无人机随机分布在边界
                side = i % 4
                if side == 0:  # 上边界
                    uav.position = (np.random.randint(0, grid_size), 10)
                elif side == 1:  # 右边界
                    uav.position = (grid_size-10, np.random.randint(0, grid_size))
                elif side == 2:  # 下边界
                    uav.position = (np.random.randint(0, grid_size), grid_size-10)
                else:  # 左边界
                    uav.position = (10, np.random.randint(0, grid_size))
            self.uavs.append(uav)
            
        # 创建多个动态目标
        self.multi_target = MultiTarget(grid_size, num_targets=2)
        # 保持向后兼容性
        self.target = self.multi_target.targets[0] if self.multi_target.targets else None
        
        # 合围参数
        self.encirclement_radius = 40  # 合围半径
        self.min_encirclement_distance = 0 # 最小合围距离
        self.max_encirclement_distance = 60  # 最大合围距离
        
        # 奖励函数参数
        self.STEP_PENALTY = -2
        self.TIME_EFFICIENCY_WEIGHT = 50
        self.COMPLETION_BASE_REWARD = 300
        self.PROGRESS_REWARD_SCALE = 150
        self.COLLISION_PENALTY = -15

    def reset(self, seed=None):
        """重置环境"""
        # 如果提供了种子，重新设置随机种子
        if seed is not None:
            np.random.seed(seed)
            random.seed(seed)
            
        # 重置多目标位置
        self.multi_target.reset()
        # 保持向后兼容性
        self.target = self.multi_target.targets[0] if self.multi_target.targets else None
        
        # 重置无人机位置
        for i, uav in enumerate(self.uavs):
            if i == 0:
                uav.position = (10, 10)
            elif i == 1:
                uav.position = (self.grid_size-10, 10)
            elif i == 2:
                uav.position = (10, self.grid_size-10)
            elif i == 3:
                uav.position = (self.grid_size-10, self.grid_size-10)
            elif i == 4:
                uav.position = (self.grid_size//2, 10)
            else:
                side = i % 4
                if side == 0:
                    uav.position = (np.random.randint(0, self.grid_size), 10)
                elif side == 1:
                    uav.position = (self.grid_size-10, np.random.randint(0, self.grid_size))
                elif side == 2:
                    uav.position = (np.random.randint(0, self.grid_size), self.grid_size-10)
                else:
                    uav.position = (10, np.random.randint(0, self.grid_size))
            uav.history_position = []
            uav.energy = 100
            
    def check_encirclement(self):
        """检查是否成功合围任一目标"""
        # 检查每个目标是否被合围
        for target in self.multi_target.targets:
            target_pos = np.array(target['position'])
            uav_positions = [np.array(uav.position) for uav in self.uavs]
            
            # 检查距离
            distances = [np.linalg.norm(pos - target_pos) for pos in uav_positions]
            in_range_count = sum(1 for d in distances if d <= 50)  # 放宽距离要求
            
            # 只需要3个无人机在范围内即可合围一个目标
            if in_range_count >= 3:
                return True
        
        return False
    
    def check_target_encirclement(self, target_id):
        """检查指定目标是否被合围"""
        if target_id >= len(self.multi_target.targets):
            return False
            
        target_pos = np.array(self.multi_target.targets[target_id]['position'])
        uav_positions = [np.array(uav.position) for uav in self.uavs]
        
        distances = [np.linalg.norm(pos - target_pos) for pos in uav_positions]
        in_range_count = sum(1 for d in distances if d <= 50)
        
        return in_range_count >= 3

    def calculate_encirclement_progress(self):
        """计算合围进度 (0-1) - 考虑所有目标的最佳进度"""
        max_progress = 0
        
        # 计算每个目标的合围进度，取最大值
        for target in self.multi_target.targets:
            target_pos = np.array(target['position'])
            uavs_in_range = 0
            
            for uav in self.uavs:
                distance = np.linalg.norm(np.array(uav.position) - target_pos)
                if distance <= self.max_encirclement_distance:
                    uavs_in_range += 1
            
            progress = uavs_in_range / len(self.uavs)
            max_progress = max(max_progress, progress)
        
        return max_progress
    
    def get_dynamic_weights(self, episode, total_episodes=300):
        """稳定的动态权重策略"""
        progress = episode / total_episodes
        
        # 减小权重变化幅度，保持训练稳定性
        if progress < 0.4:
            return {
                'distance': 1.0,
                'efficiency': 0.2,
                'completion': 1.5,
                'progress': 1.0
            }
        elif progress < 0.8:
            return {
                'distance': 1.0, 
                'efficiency': 0.25,
                'completion': 1.6,
                'progress': 1.1
            }
        else:
            return {
                'distance': 1.0,
                'efficiency': 0.3,
                'completion': 1.7,
                'progress': 1.2
            }
    
    def calculate_optimized_reward(self, uav, current_step, max_steps, episode=0):
        """稳定的奖励函数 - 适配多目标"""
        uav_pos = np.array(uav.position)
        
        # 找到最近的目标
        min_distance = float('inf')
        closest_target_pos = None
        for target in self.multi_target.targets:
            target_pos = np.array(target['position'])
            distance = np.linalg.norm(uav_pos - target_pos)
            if distance < min_distance:
                min_distance = distance
                closest_target_pos = target_pos
        
        # 获取动态权重
        weights = self.get_dynamic_weights(episode)
        
        # 1. 降低时间效率奖励的数值范围
        if self.check_encirclement():
            efficiency_ratio = (max_steps - current_step) / max_steps
            time_efficiency_bonus = 10 + efficiency_ratio * 15  # 降低到10-25
        else:
            progress_ratio = self.calculate_encirclement_progress()
            time_efficiency_bonus = progress_ratio * 8  # 降低基础奖励
        
        time_efficiency_bonus *= weights['efficiency']
        
        # 2. 降低分层距离奖励的数值 - 使用最近目标
        distance_to_target = min_distance
        if distance_to_target <= self.min_encirclement_distance:
            distance_reward = 30  # 从120降低到30
        elif distance_to_target <= self.max_encirclement_distance:
            distance_reward = 20  # 从90降低到20
        elif distance_to_target <= 100:
            distance_reward = 10  # 从50降低到10
        else:
            distance_reward = 5   # 从15降低到5
        
        distance_reward *= weights['distance']
        
        # 3. 降低进度奖励规模
        encirclement_progress = self.calculate_encirclement_progress()
        progress_reward = encirclement_progress * 20 * weights['progress']  # 从150降低到20
        
        # 4. 重新设计完成奖励
        if self.check_encirclement():
            # 大幅提高完成奖励，使其成为主要激励
            completion_bonus = 100 * weights['completion']  # 固定高额完成奖励
        else:
            completion_bonus = 0
        
        # 5. 简化步长惩罚
        step_penalty = -0.5  # 固定小额惩罚
        
        # 6. 避碰惩罚保持不变
        collision_penalty = 0
        for other_uav in self.uavs:
            if other_uav != uav:
                other_pos = np.array(other_uav.position)
                if np.linalg.norm(uav_pos - other_pos) < 20:
                    collision_penalty += self.COLLISION_PENALTY
        
        # 7. 简化协作奖励
        cooperation_reward = 0
        if len(self.uavs) > 1:
            angles = []
            for other_uav in self.uavs:
                if other_uav != uav:
                    other_pos = np.array(other_uav.position)
                    other_distance = np.linalg.norm(other_pos - target_pos)
                    if other_distance <= self.max_encirclement_distance:
                        angles.append(1)  # 简化角度计算
            
            if len(angles) > 0:
                cooperation_reward = 5 * len(angles)  # 从20降低到5
        
        # 计算总奖励
        total_reward = (distance_reward + progress_reward + completion_bonus + 
                       time_efficiency_bonus + cooperation_reward + 
                       step_penalty + collision_penalty)
        
        # 添加奖励范围限制，防止异常值
        total_reward = np.clip(total_reward, -50, 200)
        
        return total_reward
    
    def calculate_simplified_reward(self, uav):
        """保留原始简化奖励函数作为备选 - 适配多目标"""
        uav_pos = np.array(uav.position)
        
        # 找到最近的目标
        min_distance = float('inf')
        for target in self.multi_target.targets:
            target_pos = np.array(target['position'])
            distance = np.linalg.norm(uav_pos - target_pos)
            min_distance = min(min_distance, distance)
        
        # 1. 主要奖励：接近最近的目标
        distance_to_target = min_distance
        max_distance = np.sqrt(2) * self.grid_size
        approach_reward = 50 * (1 - distance_to_target / max_distance)
        
        # 2. 合围成功奖励
        encirclement_reward = 200 if self.check_encirclement() else 0
        
        # 3. 简单的避碰惩罚
        collision_penalty = 0
        for other_uav in self.uavs:
            if other_uav != uav:
                other_pos = np.array(other_uav.position)
                if np.linalg.norm(uav_pos - other_pos) < 20:
                    collision_penalty -= 10
        
        return approach_reward + encirclement_reward + collision_penalty
    
    def perform_action(self, uav, cruise_action, current_step=0, max_steps=150, episode=0, use_optimized=True):
        """执行动作并返回奖励"""
        # 执行无人机动作
        movement_reward, _ = uav.execute_cruise_action(cruise_action)
        
        # 移动所有目标（动态目标）
        self.multi_target.move_all_targets()
        
        # 检查是否有目标需要逃脱
        uav_positions = [uav.position for uav in self.uavs]
        for i, target in enumerate(self.multi_target.targets):
            # 如果目标周围有3个或更多无人机，尝试逃脱
            target_pos = np.array(target['position'])
            nearby_uavs = sum(1 for pos in uav_positions 
                            if np.linalg.norm(np.array(pos) - target_pos) <= 60)
            if nearby_uavs >= 3:
                self.multi_target.try_escape(i, uav_positions)
        
        # 选择奖励函数
        if use_optimized:
            encirclement_reward = self.calculate_optimized_reward(uav, current_step, max_steps, episode)
        else:
            encirclement_reward = self.calculate_simplified_reward(uav)
        
        total_reward = movement_reward + encirclement_reward
        return total_reward
    
    def get_global_state(self):
        """获取全局状态 - 包含所有目标"""
        states = []
        for uav in self.uavs:
            states.extend(uav.get_state())
        # 添加所有目标的状态
        states.extend(self.multi_target.get_state())
        return np.array(states)
    
    def is_done(self):
        """检查任务是否完成"""
        return self.check_encirclement()